package com.zoe.tree.util;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.Assert;
import cn.hutool.core.util.ObjectUtil;

import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.BiConsumer;
import java.util.function.Function;

/**
 * Set<TreeBeanComparable> buildTree = Tree.<TreeBeanComparable, Long>builder()
 * <br/>                                         .collection(dataSetList)
 * <br/>                                         .nodeId(TreeBeanComparable::getId)
 * <br/>                                         .pidId(TreeBeanComparable::getPid)
 * <br/>                                         .setChildren(TreeBeanComparable::addChildren)
 * <br/>                                         .build();
 * @author zbg
 * @date 2023/08/18
 * @description
 */
public class SortTreeBuild<E extends Comparable<E>, K> {
    private Collection<E> collection;
    private Function<E, K> nodeIdFun;
    private Function<E, K> pidFun;
    private BiConsumer<E, E> addChildren;

    public static <E extends Comparable<E>, K> SortTreeBuild<E, K> builder() {
        return new SortTreeBuild<>();
    }

    public SortTreeBuild<E, K> collection(Collection<E> collection) {
        this.collection = collection;
        return this;
    }

    public SortTreeBuild<E, K> nodeId(Function<E, K> nodeIdFun) {
        this.nodeIdFun = nodeIdFun;
        return this;
    }

    public SortTreeBuild<E, K> pidId(Function<E, K> pidFun) {
        this.pidFun = pidFun;
        return this;
    }

    public SortTreeBuild<E, K> setChildren(BiConsumer<E, E> addChildren) {
        this.addChildren = addChildren;
        return this;
    }

    public Set<E> build() {
        if (CollUtil.isEmpty(collection)) {
            return new TreeSet<>();
        }
        Assert.isTrue(ObjectUtil.isAllNotEmpty(nodeIdFun, pidFun, addChildren), "树构建失败,确认\n" +
                "【nodeIdFun pidFun addChildren 】属性已被初始化");

        return geteTreeSet();
    }


    private TreeSet<E> geteTreeSet() {
        TreeSet<E> treeSet = new TreeSet<>();
        Map<K, E> listToMap = LambdaUtils.listToMap(collection, nodeIdFun);
        for (E e : collection) {
            K nodeId = pidFun.apply(e);
            E parent = listToMap.get(nodeId);
            if (ObjectUtil.isNull(parent)) {
                treeSet.add(e);
            } else {
                addChildren.accept(parent, e);
            }
        }
        return treeSet;
    }

}
